#!/usr/bin/env python3
import shutil
import os 
from email.mime.text import MIMEText
import smtplib

from config import config

def check_disk_usage(path="/data",threshold=70):
        """
        检查指定路径所在磁盘的使用百分比
        :param path: 指定路径
        :param threshold: 使用率阀值: (百分比)
        :return: 是否超过阀值
        """
        total, used, free = shutil.disk_usage(path)

        percent_used = (used / total) * 100
        print(f"[+] 磁盘 {path} 使用率: {percent_used:.2f}%")

        return percent_used > threshold

#控制台报警

def alert_console(threshold):
        print(f"[-] 警告: 磁盘使用率已超过 {threshold}%! ")

# 发送邮件告警 (需要配置smtp)
def alert_email(subject, body, sender, receiver, smtp_server, smtp_port, smtp_user, smtp_pass):
        msg = MIMEText(body)
        msg['Subject'] = subject
        msg['From'] = sender
        msg['To'] = receiver

        try:
                with smtplib.SMTP(smtp_server, smtp_port) as server:
                        server.starttls()
                        server.login(smtp_user, smtp_pass)
                        server.sendmail(sender, receiver, msg.as_string())
                print("[+] 邮件报警发送成功")
        except Exception as e:
                print("[-] 邮件发送失败: ", str(e))

#主函数
def main():
        path = config.get("path", "/data")
        threshold = config.get("threshold", 70)
        enable_email = config.get("enable_email", False)

        if check_disk_usage(path, threshold):
                alert_console(threshold)

                if enable_email:
                        smtp_server = config.get("smtp_server")
                        smtp_port = config.get("smtp_port")
                        smtp_user = config.get("smtp_user")
                        smtp_pass = config.get("smtp_pass")
                        sender = config.get("sender")
                        receiver = config.get("receiver")

                        if not all([smtp_server, smtp_port, smtp_user, smtp_pass, sender, receiver]):
                                print("[-] 邮件报警配置不完整")
                                return
                        subject = f"磁盘空间告警: {path}"
                        total, used, _=shutil.disk_usage(path)
                        precent_used = (used / total) * 100
                        body=f"磁盘 {path} 的使用率超过 {threshold}%! 当前使用率: {precent_used:.2f}%"

                        alert_email(
                                        subject,
                                        body,
                                        sender,
                                        receiver,
                                        smtp_server,
                                        smtp_port,
                                        smtp_user,
                                        smtp_pass
                                        )

if __name__ == "__main__":
        main()